/* Copyright (C) 2005, 2006, 2007, 2008, 2009, 2010 Mario Orsi
   This file is part of Brahms.
   Brahms is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by
   the Free Software Foundation, either version 3 of the License, or (at your option) any later version.
   Brahms is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more details.
   You should have received a copy of the GNU General Public License along with Brahms.  If not, see <http://www.gnu.org/licenses/>. */

/***************************************************************************************
 * stepCalc.c -- functions for calculating the main properties of the simulated system *
 ***************************************************************************************/

//#define DEBUG

#include "dataStructs.h"

// global variables:
OrdParSolute ordParSolute; // P2 order parameters of solute molecules
Prop headGroupTilt; /* headgroup tilt angle */ 
Prop ordParTop; /* P2 order parameter for the top-tail (saturated) bond: TAIL_A2 -> TAIL_A1 */
Prop ordParMidUnsat; /* P2 order parameter for the mid-tail unsaturated bond: TAIL_A3 -> TAIL_A2 */
Prop ordParMidSat; /* P2 order parameter for the mid-tail saturated bond: TAIL_A4 -> TAIL_A3 */
Prop ordParEnd; /* P2 order parameter for the end-tail (saturated) bond: TAIL_A5 -> TAIL_A4 */
Prop ordParTail; /* order parameter for the entire tail - simply the average of the four defined above */

VecR *currWatCoords; /* water coordinates at the current step */
VecR *prevWatCoords; /* water coordinates at the previous step */
VecR *refWatCoords;

extern const VecR *siteCoords, *siteVelocity, *siteAngMom; 
extern const RMat *siteRotMat; 
extern const int DOFs, nLipids, nSites, nSolutes, nWaters, stepAvg, zConstraint;
extern const real cholCharge, deltaT, *mass, *inertia, timeNow;
extern const Site *site;  
extern const VecR region, solInertia;
extern Prop systemTemperature, lipTemperature, watTemperature, solTemperature;
extern Prop pressure, dens, boxVolume, xyArea, lipVolume, lipArea, surfaceTension;
extern Prop totEnergy, potEnergy, rotWatTemp, transWatTemp, rotLipTemp, transLipTemp, zConstraintForce;
extern Prop headGroupDipole, ordParTop, ordParMidUnsat, ordParMidSat, ordParEnd, lowerGlyTilt, upperGlyTilt; 

/* external function declarations: */
void GetOrientPrincAxisX_FromRotMat( VecR*, RMat );
void GetOrientPrincAxisY_FromRotMat( VecR*, RMat );
void GetOrientPrincAxisZ_FromRotMat( VecR*, RMat );

/****************************************************************************************************************
 * AccumProps -- collects the measurement results and evaluates means and standard deviations upon request. 
 * Also, initialises the measurements when called from SetUpJob(). The macro employed are defined in 
 * "definitions.h". Adapted from [Rapaport, p.32]. 
 ***************************************************************************************************************/
void AccumProps( int icode )
{
  if ( icode == 0 ) { /* initialise the accumulated sums */
    PropZero( totEnergy );
    PropZero( potEnergy );
    PropZero( pressure );
    PropZero( systemTemperature ); 
    PropZero( lipTemperature ); 
    PropZero( watTemperature );
    PropZero( solTemperature );
    PropZero( rotWatTemp );
    PropZero( transWatTemp );
    PropZero( rotLipTemp );
    PropZero( transLipTemp );
    PropZero( dens );
    PropZero( boxVolume );
    PropZero( xyArea );
    PropZero( lipVolume );
    PropZero( lipArea );
    PropZero( surfaceTension );
    PropZero( headGroupDipole );
    PropZero( headGroupTilt );
    PropZero( ordParTop );
    PropZero( ordParMidUnsat );
    PropZero( ordParMidSat );
    PropZero( ordParEnd );
    PropZero( ordParTail );
    PropZero( ordParSolute.p2x );
    PropZero( ordParSolute.p2y );
    PropZero( ordParSolute.p2z );
    PropZero( zConstraintForce );
  } else if ( icode == 1 ) { /* accumulate the current values */
    PropAccum( totEnergy );
    PropAccum( potEnergy );
    PropAccum( pressure );
    PropAccum( systemTemperature );
    PropAccum( lipTemperature );
    PropAccum( watTemperature );
    PropAccum( solTemperature );
    PropAccum( rotWatTemp );
    PropAccum( transWatTemp );
    PropAccum( rotLipTemp );
    PropAccum( transLipTemp );
    PropAccum( dens );
    PropAccum( boxVolume );
    PropAccum( xyArea );
    PropAccum( lipVolume );
    PropAccum( lipArea );
    PropAccum( surfaceTension );
    PropAccum( headGroupDipole );
    PropAccum( headGroupTilt );
    PropAccum( ordParTop );
    PropAccum( ordParMidUnsat );
    PropAccum( ordParMidSat );
    PropAccum( ordParEnd );
    PropAccum( ordParTail );
    PropAccum( ordParSolute.p2x );
    PropAccum( ordParSolute.p2y );
    PropAccum( ordParSolute.p2z );
    PropAccum( zConstraintForce );
  } else if ( icode == 2 ) { /* produce the final average estimates */
    PropAvg( totEnergy, stepAvg ); /* accumulated values overwritten! */
    PropAvg( potEnergy, stepAvg );
    PropAvg( pressure, stepAvg );
    PropAvg( systemTemperature, stepAvg );
    PropAvg( lipTemperature, stepAvg );
    PropAvg( watTemperature, stepAvg );
    PropAvg( solTemperature, stepAvg );
    PropAvg( rotWatTemp, stepAvg );
    PropAvg( transWatTemp, stepAvg );
    PropAvg( rotLipTemp, stepAvg );
    PropAvg( transLipTemp, stepAvg );
    PropAvg( dens, stepAvg );
    PropAvg( boxVolume, stepAvg );
    PropAvg( xyArea, stepAvg );
    PropAvg( lipVolume, stepAvg );
    PropAvg( lipArea, stepAvg );
    PropAvg( surfaceTension, stepAvg );
    PropAvg( headGroupDipole, stepAvg );
    PropAvg( headGroupTilt, stepAvg );
    PropAvg( ordParTop, stepAvg );
    PropAvg( ordParMidUnsat, stepAvg );
    PropAvg( ordParMidSat, stepAvg );
    PropAvg( ordParEnd, stepAvg );
    PropAvg( ordParTail, stepAvg );
    PropAvg( ordParSolute.p2x, stepAvg );
    PropAvg( ordParSolute.p2y, stepAvg );
    PropAvg( ordParSolute.p2z, stepAvg );
    PropAvg( zConstraintForce, stepAvg );
  }
}

/***************************************************
 * CalcPressTens -- calculates the pressure tensor *
 ***************************************************/
void CalcPressTens( RMat *pressureTensor, /* out */
		    const RMat virialTensor ) /* in */
{
  int n;
  RMat kinEnDensTensor; // kinetic energy density tensor 

  MZero( pressureTensor->u ); /* initialize */
  DO_SITE {
    VOuter( kinEnDensTensor.u, siteVelocity[ n ], siteVelocity[ n ] );
    MScale( kinEnDensTensor.u, mass[ site[ n ].type ] );
    MMAdd( pressureTensor->u, kinEnDensTensor.u ); /* accumulate kinetic contribution to the pressure */
  }
  MMAdd( pressureTensor->u, virialTensor.u ); 
  MScale( pressureTensor->u, 1. / boxVolume.val );
}

/*************************************************************************************
 * CalcLipTemperature -- calculates temperature of the lipid component of the system *
 *************************************************************************************/
static real CalcLipTemperature() 
{
  int n, lipDOFs = 0;
  real lipidTemperature, kineticEnergyLipids = 0.; 
  
  DO_SITE { 
    if ( site[ n ].inLipid ) { /* if the site belongs to a lipid molecule */
      lipDOFs += 3; /* add translational degrees of freedom */
      kineticEnergyLipids += 0.5 * VLenSq( siteVelocity[ n ] ) * mass[ site[ n ].type ]; /* accumulate translational contribution */
      if ( SRB == site[ n ].mechType ) { /* if the site is a symmetric rigid body (e.g., a site embedded with a dipole) */
	lipDOFs += 2; /* add rotational degrees of freedom for a symmetric rigid body */
	kineticEnergyLipids += 0.5 * VLenSq( siteAngMom[ n ] ) / inertia[ site[ n ].type ]; /* accumulate rotational contribution */
      }
    }
  }

  lipidTemperature = 2. * kineticEnergyLipids / ( kB_IN_BRAHMS_UNITS * lipDOFs ); 
  return ( lipidTemperature );
}

/****************************************************************
 * CalcKinEnergy() - returns total kinetic energy of the system * 
 ****************************************************************/
real CalcKinEnergy() 
{
  int n;
  real kinEnergy = 0.; // initialise
  DO_SITE {
    kinEnergy += 0.5 * VLenSq( siteVelocity[ n ] ) * mass[ site[ n ].type ];  /* accumulate contribution from translational motion */
    if ( SRB == site[ n ].mechType ) /* if the site is a symmetric rigid body (e.g., embedded with a dipole) */
      kinEnergy += 0.5 * VLenSq( siteAngMom[ n ] ) / inertia[ site[ n ].type ]; /* accumulate */
    else if ( NSRB == site[ n ].mechType ) { /* else if the site is a non-symmetric rigid bodies (e.g., SSD water or rigid solute) */
      if ( site[ n ].type == SOLUTE ) {
	kinEnergy += 0.5 * VSumSqDiv( siteAngMom[ n ], solInertia );
      } else {
	printf("Error from CalcKinEnergy -- undefined case. Exiting...\n");
	exit(0);
      }
    }
  }
  return kinEnergy;
}

static real CalcSolTemperature() // returns current temperature for solute 
{
  int n, soluteDOFs, totalSoluteDOFs = 0; // DOF = Degree Of Freedom 
  real soluteTemp, soluteKinEnergy = 0.; // initialise

  if ( zConstraint ) soluteDOFs = 5; // 1 DOF is lost due to the solute being z-constraint
  else soluteDOFs = 6; // general case

  DO_SITE {
    if ( SOLUTE == site[ n ].type ) {
      totalSoluteDOFs += soluteDOFs;
      soluteKinEnergy += VLenSq( siteVelocity[ n ] ) * mass[ site[ n ].type ];  // translational contribution
      soluteKinEnergy += VSumSqDiv( siteAngMom[ n ], solInertia );                // rotational contribution from non-symmetric rigid bodies
    }
  }
  
  soluteKinEnergy *= .5;
  soluteTemp = 2. * soluteKinEnergy / ( kB_IN_BRAHMS_UNITS * totalSoluteDOFs ); 
  return ( soluteTemp );
}

static real CalcWatTemperature() // returns current kinetic energy for lipids and waters separately
{
  int n, watDOFs = 0;
  real watTemp, kinEnergyWats = 0.; // initialise

  DO_SITE {
    if ( WATER_TYPE == site[ n ].type ) {
      watDOFs += 5; // 3 translational degrees of freedom + 2 rotational degrees of freedom
      kinEnergyWats += mass[ site[ n ].type ] * VLenSq( siteVelocity[ n ] );  // translational contribution
      kinEnergyWats += VLenSq( siteAngMom[ n ] ) / inertia[ site[ n ].type ]; // rotational contribution from symmetric rigid bodies
    }
  }
  
  kinEnergyWats *= .5;
  watTemp = 2. * kinEnergyWats / ( kB_IN_BRAHMS_UNITS * watDOFs ); 

  return ( watTemp );
}

static real ComputeHGDipole() // summing up and averaging HG dipole moments & HG tilt angle
{                      
  int n; // counter used in DO_SITE macro
  real tiltAngle = 0., choline_phosphate_distanceMagnitude, headgroupDipoleSum = 0.; // used to accumulate the hg dipole over all the lipids
  VecR choline_phosphate_distanceVector; // distance between the two point charges (chol & phos)
  VecR zPlus, zMinus;
  VSet( zPlus, 0., 0., 1.);
  VSet( zMinus, 0., 0., -1.);
  DO_SITE { 
    if ( site[ n ].lipidUnit == CHOLINE_OR_AMINE ) { 
      // it is assumed that the CHOLINE (or AMINE) sites always immediately precede PHOSPHATE sites in the site array
      VSub( choline_phosphate_distanceVector, siteCoords[ n ], siteCoords[ n + 1 ] ); //  = rChol - rPhos
      VWrapAll( choline_phosphate_distanceVector ); // apply periodic boundary conditions
      choline_phosphate_distanceMagnitude = VLen( choline_phosphate_distanceVector );
      headgroupDipoleSum += choline_phosphate_distanceMagnitude; 
      if ( siteCoords[ n + 2 ].z > 0. ) /* siteCoords[ n + 2 ] corresponds to a GLYCEROL site */
	// computing hg-topBilayer angle as cos( dipole,normal )
	tiltAngle +=  VDot( choline_phosphate_distanceVector, zPlus ) / choline_phosphate_distanceMagnitude; 
      else 
	// computing hg-bottomBilayer angle as cos( dipole,normal )
	tiltAngle +=  VDot( choline_phosphate_distanceVector, zMinus ) / choline_phosphate_distanceMagnitude; 
    }
  }
  tiltAngle /= nLipids;
  tiltAngle = 90 - ( acos( tiltAngle ) * 180. / PI ); 
  headGroupTilt.val = tiltAngle;
  headgroupDipoleSum *= cholCharge; // assuming |cholCharge| = |phosCharge|
  return ( headgroupDipoleSum / nLipids );
}

/****************************************************************************************************************
 * ComputeTailOrderParameters -- Calculates the orientational order parameters (also known as "P2") of lipid
 *                               tails. The order parameter is defined as P2 = (3*cos^2(theta)-1) / 2,
 *                               with theta the angle between the bilayer normal (conventionally the z-axis)
 *                               and a bond vector connecting two adjacent sites along the lipid hydrocarbon tail. 
 *
 *                               The calculation is performed assuming the elementary relation: 
 *                               cos(theta) = (bondVector.zAxisUnitVector)/bondVectorLength
 *
 *                               Four separate order parameters are evaulated (ordParTop, ordParMidUnsat, 
 *                               ordParMidSat and ordParEnd), together with the overall tail average (ordParTail).
 *
 *                               This function assumes that the bond vector connects the "n" site to the 
 *                               "n-1" site, hence it relies on the arbitrary topology defined in "dataStructs.h" 
 ****************************************************************************************************************/
static void ComputeTailOrderParameters() 
{
  int n;
  VecR bilayerNormal, bond;
  VSet( bilayerNormal, 0., 0., 1. );  /* assuming bilayerNormal = z-direction */

  ordParTop.val = ordParMidUnsat.val = ordParMidSat.val = ordParEnd.val = 0.; /* initialize */

  DO_SITE { 
    switch ( site[ n ].lipidUnit ) { 
    case WATER: 
    case CHOLINE_OR_AMINE: 
    case PHOSPHATE: 
    case GLYCEROL: 
    case ESTER_A: 
    case ESTER_B:
    case TAIL_A1:
    case TAIL_B1:
      break; 
    case TAIL_A2: 
    case TAIL_B2: /* calculate "top" order parameter corresponding to the TAIL_1 <- TAIL_2 bond */
      VSub( bond, siteCoords[ n-1 ], siteCoords[ n ]); /* calculate TAIL_1 <- TAIL_2 bond vector */
      ordParTop.val += ( 3. * Sqr( VDot( bond, bilayerNormal ) / VLen( bond ) ) - 1. ) / 2.;
      break;
    case TAIL_A3: 
    case TAIL_B3: /* calculate "midUnsat" order parameter corresponding to the TAIL_2 <- TAIL_3 bond */
      VSub( bond, siteCoords[ n-1 ], siteCoords[ n ]); /* calculate TAIL_2 <- TAIL_3 bond vector */
      ordParMidUnsat.val += ( 3. * Sqr( VDot( bond, bilayerNormal ) / VLen( bond ) ) - 1. ) / 2.;
      break;
    case TAIL_A4: 
    case TAIL_B4: /* calculate "midSat" order parameter corresponding to the TAIL_3 <- TAIL_4 bond */
      VSub( bond, siteCoords[ n-1 ], siteCoords[ n ]); /* calculate TAIL_3 <- TAIL_4 bond vector */
      ordParMidSat.val += ( 3. * Sqr( VDot( bond, bilayerNormal ) / VLen( bond ) ) - 1. ) / 2.;
      break;
    case TAIL_A5: 
    case TAIL_B5: /* calculate "end" order parameter corresponding to the TAIL_4 <- TAIL_5 bond */
      VSub( bond, siteCoords[ n-1 ], siteCoords[ n ]); /* calculate TAIL_4 <- TAIL_5 bond vector */
      ordParEnd.val += ( 3. * Sqr( VDot( bond, bilayerNormal ) / VLen( bond ) ) - 1. ) / 2.;
      break;
    default: printf("Warning: something wrong in the calculation of the tail order parameters.\n");
    } 
  }

  ordParTop.val /= ( nLipids * 2. ); // "*2" because there are two tail sites in each (top/midUnsat/midSat/end) segment
  ordParMidUnsat.val /= ( nLipids * 2. );
  ordParMidSat.val /= ( nLipids * 2. );
  ordParEnd.val /= ( nLipids * 2. );
  ordParTail.val = ( ordParTop.val + ordParMidUnsat.val + ordParMidSat.val + ordParEnd.val ) / 4; /* average over the whole tail */
}

static void ComputeSoluteOrderParameters() {
  int n;
  VecR bilayerNormal, xBodyAxis, yBodyAxis, zBodyAxis;
  VSet( bilayerNormal, 0., 0., 1. );  // assuming bilayerNormal = z-direction
  
  ordParSolute.p2x.val = ordParSolute.p2y.val = ordParSolute.p2z.val = 0.; // initialise

  for ( n = 1; n <= nSolutes; n++ ) { // assuming the solutes are the last nSolutes elements of the site array 
    GetOrientPrincAxisX_FromRotMat( &xBodyAxis, siteRotMat[ nSites - n ] ); 
    GetOrientPrincAxisY_FromRotMat( &yBodyAxis, siteRotMat[ nSites - n ] ); 
    GetOrientPrincAxisZ_FromRotMat( &zBodyAxis, siteRotMat[ nSites - n ] ); 
    //    printf("% f % f % f\n", xBodyAxis.x, xBodyAxis.y, xBodyAxis.z );
    ordParSolute.p2x.val += ( 3. * Sqr( VDot( xBodyAxis, bilayerNormal ) ) - 1. ) / 2.; 
    ordParSolute.p2y.val += ( 3. * Sqr( VDot( yBodyAxis, bilayerNormal ) ) - 1. ) / 2.;
    ordParSolute.p2z.val += ( 3. * Sqr( VDot( zBodyAxis, bilayerNormal ) ) - 1. ) / 2.;
  }
  
  ordParSolute.p2x.val /= nSolutes;
  ordParSolute.p2y.val /= nSolutes;
  ordParSolute.p2z.val /= nSolutes;
}

/************************************************************************************************************
 * CountPossiblePermeatingWaters -- keeps track of water motion along the direction normal to the bilayer   *
 *                                  plane (conventionally the z-direction) and detects possible full penet- *
 *                                  rations through the membrane (water "crossing events")                  *
 *                                  This function assumes that the origin of the frame of reference is in   *
 *                                  the centre of the simulation region.                                    *				    
 ************************************************************************************************************/
void CountPossiblePermeatingWaters() 
{
  int n; /* counter */
  const real MEMBRANE_THICKNESS = 3; /* [nm] reference membrane thickness, used to identify a permeation event. 
				      A water molecule crossing event is counted when the molecule covers 
				      an intramembrane  distance corresponding to MEMBRANE_THICKNESS along the 
				      direction normal to the bilayer (typically, the z-direction). The initial
				      posistion is also taken into account (see code below ). */
  real zDistanceFromPreviousTimestep, zDistanceFromReference, refDistFromHydroCore;
  FILE *fPtr;

  DO_WATER {
    zDistanceFromPreviousTimestep = fabs( currWatCoords[ n ].z - prevWatCoords[ n ].z );
    if ( zDistanceFromPreviousTimestep > 0.5 * region.z ) { // this implies that the water site has crossed the periodic boundary 
      refWatCoords[ n ].z =  currWatCoords[ n ].z;  // in which case, just reset reference z-coordinate
    } else { 
      zDistanceFromReference = fabs( currWatCoords[n].z - refWatCoords[n].z ); /* compute z-displacement so far */
      refDistFromHydroCore = fabs( refWatCoords[ n ].z ) - MEMBRANE_THICKNESS / 2.;
      if ( ( zDistanceFromReference > ( MEMBRANE_THICKNESS + refDistFromHydroCore ) ) && /* if the site has crossed the hydrophobic core */
	   ( ( refWatCoords[ n ].z * currWatCoords[ n ].z ) < 0 ) ) { /* current and reference cooordinates must have opposite signs */ 
	fPtr = fopen( "waterPermeationEvents.log", "a" ); 
	fprintf( fPtr, "%7.3f ns: Water #%d has crossed the bilayer, travelling from ", timeNow*TIME_ns, n );
	fprintf( fPtr, "( % 4.2f % 4.2f % 4.2f ) nm to ", refWatCoords[ n ].x, refWatCoords[ n ].y, refWatCoords[ n ].z );
	fprintf( fPtr, "( % 4.2f % 4.2f % 4.2f ) nm.\n", currWatCoords[ n ].x, currWatCoords[ n ].y, currWatCoords[ n ].z ); 
	fclose( fPtr );
	refWatCoords[ n ] = currWatCoords[ n ]; // reset reference water location
      }
    }
  }
}

/*************************************************************************************
 * EvalProps -- Evaluates properties of interest. Adapted from [Rapaport, p. 31, 56] *
 *************************************************************************************/
void EvalProps( RMat virialTensor, RMat *pressureTensor ) 
{
  real systemKineticEnergy;
  const real WATER_MOLECULAR_MASS = 18; /* mass of one water molecule in atomic mass units (Brahms units by definition) */
  const real WATER_MOLECULAR_VOLUME = 0.03; /* volume of one water molecule in nm^3 (Brahms units by definition )*/

  lipTemperature.val = CalcLipTemperature();
  watTemperature.val = CalcWatTemperature();
  solTemperature.val = CalcSolTemperature(); 
  
  systemKineticEnergy = CalcKinEnergy();
  systemTemperature.val = 2. * systemKineticEnergy / ( kB_IN_BRAHMS_UNITS * DOFs ); // DOFs calculated in startUp.c
  potEnergy.val /= nSites;
  totEnergy.val = systemKineticEnergy / nSites + potEnergy.val;

  if ( nSites==nWaters ) dens.val = nSites * WATER_MOLECULAR_MASS / boxVolume.val; 
  else dens.val = nSites / boxVolume.val; 
  
  lipVolume.val = ( boxVolume.val - nWaters * WATER_MOLECULAR_VOLUME ) / nLipids; 
  lipArea.val = xyArea.val / ( nLipids / 2. ); // printf("%f\n", lipArea.val*LENGTH_Angstrom*LENGTH_Angstrom);

  CalcPressTens( pressureTensor, virialTensor );  // evaluating pressure tensor

  pressure.val =  Trace( pressureTensor->u ) / 3.; // printf("pressure.val= % f atm (pTensor)\n\n", pressure.val*PRESSURE_atm );

  surfaceTension.val = 0.5 * region.z * ( pressureTensor->u[ 8 ] - 0.5 * ( pressureTensor->u[ 0 ] + pressureTensor->u[ 4 ] ) ); // per interface!
  
  headGroupDipole.val = ComputeHGDipole(); // also headGroupTilt computed in this function
  ComputeTailOrderParameters();
  if ( zConstraint ) ComputeSoluteOrderParameters();
}
